Your classifier is secretly an energy based model and you should treat it like one
一般的な学習
あるNNを$ f_\theta(x)とすると, 出力の$ y番目を$ f_\theta(x)[y\rbrackとして, softmaxは以下のように表される $ p_{\theta}(y|{\bf x}) = \frac{\exp{\left(f_{\theta}({\bf x})[y\rbrack \right)} } { \sum_{y^{\prime}}\exp{\left(f_{\theta}({\bf x})[y^{\prime}\rbrack \right)} }
$ p_{\theta}(\boldsymbol{x},y) = \frac{\textrm{exp}(-E_{\theta}(\boldsymbol{x},y))}{Z_{\theta}}
と定義されるので, エネルギー関数 $ E_{\theta}(\boldsymbol{x},y)を
$ E_{\theta}(\boldsymbol{x},y) := -\ln(f_{\theta}({\boldsymbol{x}})[y\rbrack)
と定義すれば, エネルギー関数 $ E_{\theta}(\boldsymbol{x})は$ yについて周辺化して
$ E_\theta(\boldsymbol{x}) := -\sum_y \ln(f_\theta(\boldsymbol{x})[y\rbrack)
これをJoint Energy-Based Model (JEM)と呼ぶ
図にすると下のような感じ
https://gyazo.com/cb2e04dc6fbbdb969522c4df27d6647c
このとき, 最適化したい対数尤度は
$ \ln(p_\theta(y|x)) = \ln(p_\theta(x)) + \ln(p_\theta(x,y))
第二項はそのままクロスエントロピーとして最適化すればよいので第二項だけ考える
$ \nabla_{\theta}E_{\theta}(x_{train}) - \mathbb{E}_{sample}[\nabla_{\theta}E_{\theta}(x_{sample})\rbrack